use super::*;
use core::mem::{size_of_val, MaybeUninit};
use core::marker::PhantomData;
use hierr::Error;

pub struct ElfFd {
    fd: i32,
    addr: *mut libc::c_void,
    len: usize,
}

impl ElfFd {
    pub fn open(filename: &str) -> Result<Self, Error> {
        let mut buf = [0_u8; 256];
        if filename.len() >= size_of_val(&buf) {
            return Err(Error::new(-1));
        }
        unsafe {
            core::ptr::copy_nonoverlapping(
                filename.as_bytes().as_ptr(),
                buf.as_mut_ptr(),
                filename.len(),
            )
        };

        let fd = unsafe { libc::open(buf.as_ptr().cast::<i8>(), libc::O_RDONLY) };
        if fd == -1 {
            return Err(Error::last_error());
        }

        let mut fd = ElfFd {
            fd,
            addr: core::ptr::null_mut(),
            len: 0,
        };

        let mut stat = MaybeUninit::<libc::stat>::uninit();
        let ret = unsafe { libc::fstat(fd.fd, stat.as_mut_ptr()) };
        if ret == -1 {
            return Err(Error::last_error());
        }

        let len = unsafe { stat.assume_init_ref().st_size as usize };
        let addr = unsafe {
            libc::mmap(
                core::ptr::null_mut(),
                len,
                libc::PROT_READ | libc::PROT_WRITE,
                libc::MAP_PRIVATE,
                fd.fd,
                0,
            )
        };
        if addr == libc::MAP_FAILED {
            return Err(Error::last_error());
        }

        fd.addr = addr;
        fd.len = len;

        if !fd.is_elf() {
            return Err(Error::new(libc::EINVAL));
        }

        Ok(fd)
    }

    pub fn size(&self) -> usize {
        self.len
    }

    pub fn read<T: Sized>(&self, offset: usize) -> Option<&T> {
        if offset + core::mem::size_of::<T>() <= self.len {
            let addr = unsafe { self.addr.add(offset) }.cast::<T>();
            return Some(unsafe { &*addr });
        }
        None
    }

    pub fn read_mut<T: Sized>(&mut self, offset: usize) -> Option<&mut T> {
        if offset + core::mem::size_of::<T>() <= self.len {
            let addr = unsafe { self.addr.add(offset) }.cast::<T>();
            return Some(unsafe { &mut *addr });
        }
        None
    }

    pub fn read_slice<T: Sized>(&self, offset: usize, len: usize) -> Option<&[T]> {
        if offset + len * core::mem::size_of::<T>() <= self.len {
            let addr = unsafe { self.addr.add(offset) }.cast::<T>();
            return Some(unsafe { core::slice::from_raw_parts(addr, len) });
        }
        None
    }

    pub fn read_slice_mut<T: Sized>(&mut self, offset: usize, len: usize) -> Option<&mut [T]> {
        if offset + len * core::mem::size_of::<T>() <= self.len {
            let addr = unsafe { self.addr.add(offset) }.cast::<T>();
            return Some(unsafe { core::slice::from_raw_parts_mut(addr, len) });
        }
        None
    }

    pub fn read_cstr(&self, offset: usize) -> Option<&str> {
        if offset < self.len {
            let addr = unsafe { self.addr.add(offset) }.cast::<u8>();
            for n in 0..self.len - offset {
                if unsafe { *addr.add(n) } == 0 {
                    let bytes = unsafe { core::slice::from_raw_parts(addr, n) };
                    return Some(unsafe { core::str::from_utf8_unchecked(bytes) });
                }
            }
        }
        None
    }

    fn segment_offset<T: ElfBits>(&self, index: usize) -> Option<usize> {
        if !self.is_bits::<T>() {
            return None;
        }
        let ehdr = self.map(self.read::<T::ElfHdr>(0)?);
        if index >= ehdr.e_phnum() as usize {
            return None;
        }
        Some(index * ehdr.e_phentsize() as usize + ehdr.e_phoff() as usize)
    }

    pub fn read_segment<T: ElfBits>(&self, index: usize) -> Option<&T::SegmentHdr> {
        let offset = self.segment_offset::<T>(index)?;
        self.read::<T::SegmentHdr>(offset)
    }

    pub fn read_segment_mut<T: ElfBits>(&mut self, index: usize) -> Option<&mut T::SegmentHdr> {
        let offset = self.segment_offset::<T>(index)?;
        self.read_mut::<T::SegmentHdr>(offset)
    }

    pub fn read_section<T: ElfBits>(&self, index: usize) -> Option<&T::SectionHdr> {
        let offset = self.section_offset::<T>(index)?;
        self.read::<T::SectionHdr>(offset)
    }

    pub fn read_section_mut<T: ElfBits>(&mut self, index: usize) -> Option<&mut T::SectionHdr> {
        let offset = self.section_offset::<T>(index)?;
        self.read_mut::<T::SectionHdr>(offset)
    }

    fn section_offset<T: ElfBits>(&self, index: usize) -> Option<usize> {
        if !self.is_bits::<T>() {
            return None;
        }
        let ehdr = self.map(self.read::<T::ElfHdr>(0)?);
        if index >= ehdr.e_shnum() as usize {
            return None;
        }
        Some(index * ehdr.e_shentsize() as usize + ehdr.e_shoff() as usize)
    }

    pub fn read_section_name<T: ElfBits>(&self, name: u32) -> Option<&str> {
        if !self.is_bits::<T>() {
            return None;
        }
        let ehdr = self.map(self.read::<T::ElfHdr>(0)?);
        let strndx = ehdr.e_shstrndx() as usize;
        if strndx == SHN_UNDEF as usize {
            return None;
        }
        let hdr = self.read_section::<T>(strndx)?;
        let offset = hdr.sh_offset() as usize + name as usize;
        self.read_cstr(offset)
    }

    pub fn section_name(&self, name: u32) -> Option<&str> {
        if self.is_bits::<u64>() {
            self.read_section_name::<u64>(name)
        } else {
            self.read_section_name::<u32>(name)
        }
    }

    pub fn is_elf(&self) -> bool {
        self.read_slice::<u8>(0, 4) == Some(&[ELFMAG0, ELFMAG1, ELFMAG2, ELFMAG3])
    }

    pub fn host_endian(&self) -> bool {
        #[cfg(target_endian = "little")]
        const ENDIAN: u8 = ELFDATA2LSB;
        #[cfg(target_endian = "big")]
        const ENDIAN: u8 = ELFDATA2MSB;
        let endian = self.read::<u8>(EI_DATA);
        endian == Some(&ENDIAN)
    }

    pub fn map<T: Endian + Clone>(&self, val: &T) -> T {
        val.map(self.host_endian())
    }

    pub fn is_bits<T: ElfBits>(&self) -> bool {
        self.read::<u8>(EI_CLASS) == Some(&T::CLASS)
    }

    pub fn segment_iter<'a, T: ElfBits + 'a>(&'a self) -> impl Iterator<Item = T::SegmentHdr> + 'a {
        SegmentIter::<'a, T>::new(self)
    }

    pub fn segment_iter_mut<'a, T: ElfBits + 'a>(&'a mut self) -> impl Iterator<Item = &'a mut T::SegmentHdr> {
        SegmentIterMut::<'a, T>::new(self)
    }

    pub fn section_iter<'a, T: ElfBits + 'a>(&'a self) -> impl Iterator<Item = T::SectionHdr> + 'a {
        SectionIter::<'a, T>::new(self)
    }
    pub fn section_iter_mut<'a, T: ElfBits + 'a>(&'a mut self) -> impl Iterator<Item = &'a mut T::SectionHdr> {
        SectionIterMut::<'a, T>::new(self)
    }

    pub fn symbol_iter<'a, T: ElfBits + 'a>(&'a self) -> impl Iterator<Item = T::Symbol> + 'a {
        ItemIter::<'a, T, T::Symbol, _>::new(self, |ty| ty == SHT_DYNSYM || ty == SHT_SYMTAB)
    }

    pub fn symbol_iter_mut<'a, T: ElfBits + 'a>(&'a mut self) -> impl Iterator<Item = &'a mut T::Symbol> {
        ItemIterMut::<'a, T, T::Symbol, _>::new(self, |ty| ty == SHT_DYNSYM || ty == SHT_SYMTAB)
    }

    pub fn rel_iter<'a, T: ElfBits + 'a>(&'a self) -> impl Iterator<Item = T::Relocation> + 'a {
        ItemIter::<'a, T, T::Relocation, _>::new(self, |ty| ty == SHT_REL)
    }

    pub fn rel_iter_mut<'a, T: ElfBits + 'a>(&'a mut self) -> impl Iterator<Item = &'a mut T::Relocation> {
        ItemIterMut::<'a, T, T::Relocation, _>::new(self, |ty| ty == SHT_REL)
    }

    pub fn rela_iter<'a, T: ElfBits + 'a>(&'a self) -> impl Iterator<Item = T::RelocationEx> + 'a {
        ItemIter::<'a, T, T::RelocationEx, _>::new(self, |ty| ty == SHT_RELA)
    }

    pub fn rela_iter_mut<'a, T: ElfBits + 'a>(&'a mut self) -> impl Iterator<Item = &'a mut T::RelocationEx> {
        ItemIterMut::<'a, T, T::RelocationEx, _>::new(self, |ty| ty == SHT_RELA)
    }

    pub fn dyn_iter<'a, T: ElfBits + 'a>(&'a self) -> impl Iterator<Item = T::Dynamic> + 'a {
        ItemIter::<'a, T, T::Dynamic, _>::new(self, |ty| ty == SHT_DYNAMIC)
    }

    pub fn dyn_iter_mut<'a, T: ElfBits + 'a>(&'a mut self) -> impl Iterator<Item = &'a mut T::Dynamic> {
        ItemIterMut::<'a, T, T::Dynamic, _>::new(self, |ty| ty == SHT_DYNAMIC)
    }
}

impl Drop for ElfFd {
    fn drop(&mut self) {
        unsafe { libc::close(self.fd) };
        if !self.addr.is_null() {
            unsafe { libc::munmap(self.addr, self.len) };
        }
    }
}

pub struct SectionIter<'a, T: ElfBits> {
    fd: &'a ElfFd,
    index: usize,
    mark: core::marker::PhantomData<&'a T>,
}

impl<'a, T: ElfBits> SectionIter<'a, T> {
    fn new(fd: &'a ElfFd) -> Self {
        Self {
            fd,
            index: 0,
            mark: core::marker::PhantomData,
        }
    }
}

impl<T: ElfBits> Iterator for SectionIter<'_, T> {
    type Item = T::SectionHdr;
    fn next(&mut self) -> Option<Self::Item> {
        let section = self.fd.read_section::<T>(self.index);
        self.index += 1;
        section.map(|s| self.fd.map(s))
    }
}

pub struct SectionIterMut<'a, T: ElfBits> {
    fd: *mut ElfFd,
    index: usize,
    mark: core::marker::PhantomData<&'a T>,
}

impl<'a, T: ElfBits> SectionIterMut<'a, T> {
    fn new(fd: &'a mut ElfFd) -> Self {
        Self {
            fd,
            index: 0,
            mark: core::marker::PhantomData,
        }
    }
}

impl<'a, T: ElfBits> Iterator for SectionIterMut<'a, T> {
    type Item = &'a mut T::SectionHdr;
    fn next(&mut self) -> Option<Self::Item> {
        self.index += 1;
        unsafe { &mut *self.fd }.read_section_mut::<T>(self.index - 1)
    }
}

pub struct SegmentIter<'a, T: ElfBits> {
    fd: &'a ElfFd,
    index: usize,
    mark: core::marker::PhantomData<&'a T>,
}

impl <'a, T: ElfBits> SegmentIter<'a, T> {
    fn new(fd: &'a ElfFd) -> Self {
        Self {
            fd, index: 0,
            mark: core::marker::PhantomData,
        }
    }
}

impl<T: ElfBits> Iterator for SegmentIter<'_, T> {
    type Item = T::SegmentHdr;
    fn next(&mut self) -> Option<Self::Item> {
        let hdr = self.fd.read_segment::<T>(self.index);
        self.index += 1;
        hdr.map(|h| self.fd.map(h))
    }
}

pub struct SegmentIterMut<'a, T: ElfBits> {
    fd: *mut ElfFd,
    index: usize,
    mark: core::marker::PhantomData<&'a T>,
}

impl <'a, T: ElfBits> SegmentIterMut<'a, T> {
    fn new(fd: &'a mut ElfFd) -> Self {
        Self {
            fd, index: 0,
            mark: core::marker::PhantomData,
        }
    }
}

impl<'a, T: ElfBits> Iterator for SegmentIterMut<'a, T> {
    type Item = &'a mut T::SegmentHdr;
    fn next(&mut self) -> Option<Self::Item> {
        self.index += 1;
        unsafe { &mut *self.fd }.read_segment_mut::<T>(self.index - 1)
    }
}

pub struct ItemIter<'a, T, U, F> {
    fd: &'a ElfFd,
    hdr_ndx: usize,
    offset: usize,
    f: F,
    marker: PhantomData<(T, U)>,
}

impl<'a, T: ElfBits, U, F: Fn(u32) -> bool> ItemIter<'a, T, U, F> {
    fn new(fd: &'a ElfFd, f: F) -> Self {
        Self {
            fd,
            hdr_ndx: 0,
            offset: 0,
            f,
            marker: PhantomData,
        }
    }

    fn hdr(&mut self) -> Option<(T::SectionHdr, usize)> {
        let host_endian = self.fd.host_endian();
        loop {
            let hdr = self.fd.read_section::<T>(self.hdr_ndx)?.map(host_endian);
            if (self.f)(hdr.sh_type()) {
                return Some((hdr, self.offset));
            }
            self.hdr_ndx += 1;
            self.offset = 0;
        }
    }
}

impl<'a, T: ElfBits, U: Endian + Clone, F: Fn(u32) -> bool> Iterator for ItemIter<'a, T, U, F> {
    type Item = U;
    fn next(&mut self) -> Option<Self::Item> {
        let (hdr, offset) = self.hdr()?;
        let end = offset + core::mem::size_of::<Self::Item>();
        if end > hdr.sh_size() as usize {
            return None;
        }
        let offset = hdr.sh_offset() as usize + offset;
        let host_endian = self.fd.host_endian();
        self.offset = end;
        Some(self.fd.read::<Self::Item>(offset)?.map(host_endian))
    }
}

pub struct ItemIterMut<'a, T, U, F> {
    fd: *mut ElfFd,
    hdr_ndx: usize,
    offset: usize,
    f: F,
    marker: PhantomData<&'a (T, U)>,
}

impl<'a, T: ElfBits, U, F: Fn(u32) -> bool> ItemIterMut<'a, T, U, F> {
    fn new(fd: &'a mut ElfFd, f: F) -> Self {
        Self {
            fd,
            hdr_ndx: 0,
            offset: 0,
            f,
            marker: PhantomData,
        }
    }

    fn hdr(&mut self) -> Option<(T::SectionHdr, usize)> {
        let fd = unsafe { &*self.fd };
        let host_endian = fd.host_endian();
        loop {
            let hdr = fd.read_section::<T>(self.hdr_ndx)?.map(host_endian);
            if (self.f)(hdr.sh_type()) {
                return Some((hdr, self.offset));
            }
            self.hdr_ndx += 1;
            self.offset = 0;
        }
    }
}

impl<'a, T: ElfBits, U: Endian + Clone, F: Fn(u32) -> bool> Iterator for ItemIterMut<'a, T, U, F> {
    type Item = &'a mut U;
    fn next(&mut self) -> Option<Self::Item> {
        let (hdr, offset) = self.hdr()?;
        let end = offset + core::mem::size_of::<U>();
        if end > hdr.sh_size() as usize {
            return None;
        }
        let offset = hdr.sh_offset() as usize + offset;
        self.offset = end;
        unsafe { &mut *self.fd }.read_mut::<U>(offset)
    }
}

